package cn.iocoder.yudao.module.data.service.text2sql;

import cn.hutool.core.collection.ListUtil;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.module.data.service.text2sql.pojo.*;
import cn.iocoder.yudao.module.data.service.text2sql.pojo.s2sql.LLMReq;
import cn.iocoder.yudao.module.data.service.text2sql.pojo.s2sql.LLMResp;
import cn.iocoder.yudao.module.data.service.text2sql.pojo.s2sql.Text2SQLExemplar;
import jakarta.servlet.http.HttpServletRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.*;
import java.util.stream.Collectors;

import static cn.iocoder.yudao.module.data.service.text2sql.OnePassSCSqlGenStrategy.INSTRUCTION;

public class ChatSimpleTest {
    static Logger log = LoggerFactory.getLogger(ChatSimpleTest.class);



    public static void main(String[] args) {
        ChatSimpleTest chatSimpleTest = new ChatSimpleTest();
        LLMResp llmResp = chatSimpleTest.llmSqlTest();

        log.info("---------可执行sql：{}" ,llmResp.getSqlOutput());

        //通过jdbc执行sql


        // translate original sql and can be test in mysql
//        chatSimpleTest.transToSql(llmResp.getSqlOutput());
    }

    /**
     * 通过自然语言生成sql
     *
     * @return
     */
    public LLMResp llmSqlTest() {
        LLMReq llmReq = new LLMReq();

        Map<String, ChatApp> chatApp = getChatApp();
        llmReq.setChatAppConfig(chatApp);
        llmReq.setSchema(getLLmSchema());

        String queryText = "生产厂用电率近一个月的趋势";
        llmReq.setQueryText(queryText);

        llmReq.setCurrentDate(DateUtil.now());
        llmReq.setDynamicExemplars(getDynamicExemplars());

        LLMChatService llmChatService = new LLMChatService();
        LLMResp llmResp = llmChatService.generate(llmReq);

        //打印生成的sql
        String sqlOutput = llmResp.getSqlOutput();
        log.info("sqlOutput: {}", sqlOutput);

        //通过jdbc执行sql
        Connection connection = null;
        Statement statement = null;
        ResultSet resultSet = null;

        try {
            connection = DriverManager.getConnection(llmReq.getSchema().getJdbcUrl()
            , llmReq.getSchema().getUserName(), llmReq.getSchema().getPassword());
            statement = connection.createStatement();
            resultSet = statement.executeQuery(sqlOutput);
            while (resultSet.next()) {
                //获取列数
                int columnCount = resultSet.getMetaData().getColumnCount();
                for (int i = 1; i <= columnCount; i++) {
                    String columnName = resultSet.getMetaData().getColumnName(i);
                    String columnDesc = resultSet.getMetaData().getColumnLabel(i);
                    String columnValue = resultSet.getString(i);
                    log.info("columnName: {},columnDesc: {},  columnValue: {}", columnName,columnDesc, columnValue);
                }
            }
            resultSet.close();
        }catch (Exception e) {

        }finally {
            //关闭结果集
            if (resultSet != null) {
                try {
                    resultSet.close();
                } catch (Exception e) {
                    log.error("close resultSet error", e);
                }
            }


            //关闭statement
            if (statement != null) {
                try {
                    statement.close();
                } catch (Exception e) {
                    log.error("close statement error", e);
                }
            }

            //关闭连接
            if (connection != null) {
                try {
                    connection.close();
                } catch (Exception e) {
                    log.error("close connection error", e);
                }
            }
        }

        return llmResp;

    }

    private List<Text2SQLExemplar> getDynamicExemplars() {
//        Text2SQLExemplar(question=昨日发电量是多少？, sideInfo=CurrentDate=[2025-04-08], dbSchema=DatabaseType=[MYSQL], DatabaseVersion=[8.0], Table=[关键指标数据集], PartitionTimeField=[时间日期 FORMAT 'yyyy-MM-dd'], PrimaryKeyField=[], Metrics=[<发电量值 COMMENT '发电量值' AGGREGATE 'SUM'>,<发电量值 COMMENT '发电量值' AGGREGATE 'SUM'>,<发电补水率 COMMENT '发电补水率' AGGREGATE 'SUM'>,<单位发电油耗 COMMENT '单位发电油耗' AGGREGATE 'SUM'>,<发电补水率 COMMENT '发电补水率' AGGREGATE 'SUM'>,<单位发电油耗 COMMENT '单位发电油耗' AGGREGATE 'SUM'>], Dimensions=[<时间日期 FORMAT 'yyyy-MM-dd' COMMENT '时间日期'>], Values=[], sql=SELECT SUM(发电量值) AS _发电量_ FROM 关键指标数据集 WHERE 时间日期 = '2025-04-07')
        Text2SQLExemplar exemplar = new Text2SQLExemplar(
                "昨日发电量是多少？",
                "CurrentDate=[2025-04-08]",
                "DatabaseType=[MYSQL], DatabaseVersion=[8.0], Table=[关键指标数据集], PartitionTimeField=[时间日期 FORMAT 'yyyy-MM-dd'], PrimaryKeyField=[], Metrics=[<发电量值 COMMENT '发电量值' AGGREGATE 'SUM'>,<发电量值 COMMENT '发电量值' AGGREGATE 'SUM'>,<发电补水率 COMMENT '发电补水率' AGGREGATE 'SUM'>,<单位发电油耗 COMMENT '单位发电油耗' AGGREGATE 'SUM'>,<发电补水率 COMMENT '发电补水率' AGGREGATE 'SUM'>,<单位发电油耗 COMMENT '单位发电油耗' AGGREGATE 'SUM'>], Dimensions=[<时间日期 FORMAT 'yyyy-MM-dd' COMMENT '时间日期'>], Values=[]",
                "SELECT SUM(发电量值) AS _发电量_ FROM 关键指标数据集 WHERE 时间日期 = '2025-04-07'"
        );

        List<Text2SQLExemplar> exemplars = new ArrayList<>();
        exemplars.add(exemplar);

        return exemplars;
    }

    private LLMReq.LLMSchema getLLmSchema() {
//        LLMReq.LLMSchema(databaseType=MYSQL, databaseVersion=8.0, dataSetId=6, dataSetName=关键指标数据集, metrics=[SchemaElement(dataSetId=6, dataSetName=关键指标数据集, model=11, id=33, name=发电量值, bizName=val1, useCnt=9, type=METRIC, alias=[], schemaValueMaps=null, relatedSchemaElements=[], defaultAgg=sum, dataFormatType=null, order=0.0, isTag=0, description=发电量值, extInfo={}, typeParams=null), SchemaElement(dataSetId=6, dataSetName=关键指标数据集, model=11, id=37, name=单位发电油耗, bizName=val5, useCnt=0, type=METRIC, alias=[], schemaValueMaps=null, relatedSchemaElements=[], defaultAgg=sum, dataFormatType=null, order=0.0, isTag=0, description=单位发电油耗, extInfo={}, typeParams=null), SchemaElement(dataSetId=6, dataSetName=关键指标数据集, model=11, id=38, name=发电补水率, bizName=val6, useCnt=0, type=METRIC, alias=[], schemaValueMaps=null, relatedSchemaElements=[], defaultAgg=sum, dataFormatType=null, order=0.0, isTag=0, description=发电补水率, extInfo={}, typeParams=null), SchemaElement(dataSetId=6, dataSetName=关键指标数据集, model=11, id=35, name=生产厂用电率, bizName=val3, useCnt=14, type=METRIC, alias=[], schemaValueMaps=null, relatedSchemaElements=[], defaultAgg=sum, dataFormatType=null, order=0.0, isTag=0, description=生产厂用电率, extInfo={}, typeParams=null), SchemaElement(dataSetId=6, dataSetName=关键指标数据集, model=11, id=36, name=综合厂用电率, bizName=val4, useCnt=4, type=METRIC, alias=[], schemaValueMaps=null, relatedSchemaElements=[], defaultAgg=sum, dataFormatType=null, order=0.0, isTag=0, description=综合厂用电率, extInfo={}, typeParams=null), SchemaElement(dataSetId=6, dataSetName=关键指标数据集, model=11, id=34, name=供电煤耗, bizName=val2, useCnt=0, type=METRIC, alias=[], schemaValueMaps=null, relatedSchemaElements=[], defaultAgg=sum, dataFormatType=null, order=0.0, isTag=0, description=供电煤耗, extInfo={}, typeParams=null)], dimensions=[SchemaElement(dataSetId=6, dataSetName=关键指标数据集, model=11, id=31, name=数据日期, bizName=btime, useCnt=3, type=DIMENSION, alias=[], schemaValueMaps=[], relatedSchemaElements=null, defaultAgg=null, dataFormatType=null, order=0.0, isTag=0, description=数据日期, extInfo={dimension_type=partition_time, time_format=yyyy-MM-dd}, typeParams=null)], values=[], partitionTime=SchemaElement(dataSetId=6, dataSetName=关键指标数据集, model=11, id=31, name=数据日期, bizName=btime, useCnt=3, type=DIMENSION, alias=[], schemaValueMaps=[], relatedSchemaElements=null, defaultAgg=null, dataFormatType=null, order=0.0, isTag=0, description=数据日期, extInfo={dimension_type=partition_time, time_format=yyyy-MM-dd}, typeParams=null), primaryKey=null)
        LLMReq.LLMSchema schema = new LLMReq.LLMSchema();
        schema.setDatabaseType("MYSQL");
        schema.setDatabaseVersion("8.0");
        schema.setDataSetId(6L);
        schema.setDataSetName("bio_indicator");
        schema.setJdbcUrl("jdbc:mysql://192.168.10.128:3306/bidemo" +
                "?useUnicode=true&characterEncoding=UTF-8&serverTimezone=UTC");
        schema.setUserName("root");
        schema.setPassword("123456");



        // 这边的指标和维度不一定有用
//        schema.setDataSetName("关键指标数据集");
//        //metrics
//        List<SchemaElement> metrics = new ArrayList<>();
//        metrics = mockMetrics();
//        schema.setMetrics(metrics);
//        //dimensions
//        List<SchemaElement> dimensions = new ArrayList<>();
//        dimensions = mockDimensions();
//        schema.setDimensions(dimensions);

        //partitionTime
//        SchemaElement partitionTime = new SchemaElement();
//        partitionTime.setDataSetId(6L);
//        partitionTime.setDataSetName("bi_indicator");
//        partitionTime.setModel(11L);
//        partitionTime.setId(31L);
//        partitionTime.setName("数据日期");
//        partitionTime.setBizName("btime");
//        partitionTime.setUseCnt(3L);
//        partitionTime.setType(SchemaElementType.DIMENSION);
//        partitionTime.setAlias(new ArrayList<>());

        return schema;

    }


    private Map<String, ChatApp>  getChatApp() {
//        "MEMORY_REVIEW" -> {ChatApp@19867} "ChatApp(name=记忆启用评估, description=通过大模型对记忆做正确性评估以决定是否启用, prompt=#Role: You are a senior data engineer experienced in writing SQL.\n#Task: Your will be provided with a user question and the SQL written by a junior engineer,please take a review and give your opinion.\n#Rules: \n1.ALWAYS follow the output format: `opinion=(POSITIVE|NEGATIVE),comment=(your comment)`.\n2.NO NEED to check date filters as the junior engineer seldom makes mistakes in this regard.\n#Question: %s\n#Schema: %s\n#SideInfo: %s\n#SQL: %s\n#Response: , enable=false, chatModelId=2, chatModelConfig=null, appModule=null)"
//        "REWRITE_MULTI_TURN" -> {ChatApp@19869} "ChatApp(name=多轮对话改写, description=通过大模型根据历史对话来改写本轮对话, prompt=#Role: You are a data product manager experienced in data requirements.#Task: Your will be provided with current and history questions asked by a user,along with their mapped schema elements(metric, dimension and value),please try understanding the semantics and rewrite a question.#Rules: 1.ALWAYS keep relevant entities, metrics, dimensions, values and date ranges.2.ONLY respond with the rewritten question.#Current Question: {{current_question}}#Current Mapped Schema: {{current_schema}}#History Question: {{history_question}}#History Mapped Schema: {{history_schema}}#History SQL: {{history_sql}}#Rewritten Question: , enable=false, chatModelId=2, chatModelConfig=null, appModule=null)"
//        "S2SQL_CORRECTOR" -> {ChatApp@19871} "ChatApp(name=语义SQL修正, description=通过大模型对解析S2SQL做二次修正, prompt=#Role: You are a senior data engineer experienced in writing SQL.\n#Task: Your will be provided with a user question and the SQL written by a junior engineer,please take a review and help correct it if necessary.\n#Rules: 1.ALWAYS specify time range using `>`,`<`,`>=`,`<=` operator.2.DO NOT calculate date range using functions.3.SQL columns and values must be mentioned in the `#Schema`.\n#Question:{{question}} #Schema:{{schema}} #InputSQL:{{sql}} #Response:, enable=false, chatModelId=2, chatModelConfig=null, appModule=null)"
//        "SMALL_TALK" -> {ChatApp@19873} "ChatApp(name=闲聊对话, description=直接将原始输入透传大模型, prompt=#Role: You are a nice person to talk to.\n#Task: Respond quickly and nicely to the user.\n#Rules: 1.ALWAYS use the same language as the `#Current Input`.\n#History Inputs: %s\n#Current Input: %s\n#Response: , enable=false, chatModelId=null, chatModelConfig=null, appModule=null)"
//        "DATA_INTERPRETER" -> {ChatApp@19875} "ChatApp(name=结果数据解读, description=通过大模型对结果数据做提炼总结, prompt=#Role: You are a data expert who communicates with business users everyday.\n#Task: Your will be provided with a question asked by a user and the relevant result data queried from the databases, please interpret the data and organize a brief answer.\n#Rules: \n1.ALWAYS respond in the use the same language as the `#Question`.\n2.ALWAYS reference some key data in the `#Answer`.\n#Question:{{question}} #Data:{{data}} #Answer:, enable=false, chatModelId=2, chatModelConfig=null, appModule=null)"
//        "S2SQL_PARSER" -> {ChatApp@19818} "ChatApp(name=语义SQL解析, description=通过大模型做语义解析生成S2SQL, prompt=#Role: You are a data analyst experienced in SQL languages.\n#Task: You will be provided with a natural language question asked by users,please convert it to a SQL query so that relevant data could be returned by executing the SQL query against underlying database.\n#Rules:\n1.SQL columns and values must be mentioned in the `Schema`, DO NOT hallucinate.\n2.ALWAYS specify time range using `>`,`<`,`>=`,`<=` operator.\n3.DO NOT include time range in the where clause if not explicitly expressed in the `Question`.\n4.DO NOT calculate date range using functions.\n5.ALWAYS use `with` statement if nested aggregation is needed.\n6.ALWAYS enclose alias declared by `AS` command in underscores.\n7.Alias created by `AS` command must be in the same language ast the `Question`.\n#Exemplars: {{exemplar}}\n#Query: Question:{{question}},Schema:{{schema}},SideInfo:{{information}}, enable=true, chatModelId=2, chatModelConfig=ChatModelConfig(provider=OPEN_AI"
//        "REWRITE_ERROR_MESSAGE" -> {ChatApp@19878} "ChatApp(name=异常提示改写, description=通过大模型将异常信息改写为更友好和引导性的提示用语, prompt=#Role: You are a data business partner who closely interacts with business people.\n#Task: Your will be provided with user input, system output and some examples, please respond shortly to teach user how to ask the right question, by using `Examples` as references.#Rules: ALWAYS respond with the same language as the `Input`.\n#Input: {{user_question}}\n#Output: {{system_message}}\n#Examples: {{examples}}\n#Response: , enable=true, chatModelId=2, chatModelConfig=ChatModelConfig(provider=OPEN_AI, baseUrl=https://dashscope.aliyuncs.com/compatible-mode/v1, apiKey=sk-e395e38b133a44b98142207c96970bb5, modelName=qwen-plus-latest, apiVersion=2024-02-01, temperature=0.0, timeOut=60, endpoint=null, secretKey=null, topP=null, maxRetries=3, logRequests=false, logResponses=false, enableSearch=false), appModule=null)"
        Map<String, ChatApp> chatAppMap = new HashMap<>();
        ChatApp memoryReview = new ChatApp();
        memoryReview.setName("记忆启用评估");
        memoryReview.setDescription("通过大模型对记忆做正确性评估以决定是否启用");
        memoryReview.setPrompt("#Role: You are a senior data engineer experienced in writing SQL.\n#Task: Your will be provided with a user question and the SQL written by a junior engineer,please take a review and give your opinion.\n#Rules: \n1.ALWAYS follow the output format: `opinion=(POSITIVE|NEGATIVE),comment=(your comment)`.\n2.NO NEED to check date filters as the junior engineer seldom makes mistakes in this regard.\n#Question: %s\n#Schema: %s\n#SideInfo: %s\n#SQL: %s\n#Response: ");
        memoryReview.setEnable(false);
        memoryReview.setChatModelId(2);
        memoryReview.setChatModelConfig(null);
        memoryReview.setAppModule(null);
        ChatModelConfig memoryChatModelConfig = new ChatModelConfig("OPEN_AI",
                "https://dashscope.aliyuncs.com/compatible-mode/v1",
                "sk-e395e38b133a44b98142207c96970bb5", "qwen-plus-latest",
                "2024-02-01", 0.0, 60L, null, null, null, 3,
                false, false, false);
        memoryReview.setChatModelConfig(memoryChatModelConfig);
        //put the memoryReview to chatAppMap
        chatAppMap.put("MEMORY_REVIEW", memoryReview);

        ChatApp s2sqlParser = new ChatApp();
        s2sqlParser.setName("语义SQL解析");
        s2sqlParser.setDescription("通过大模型做语义解析生成S2SQL");
        s2sqlParser.setPrompt(INSTRUCTION);
        s2sqlParser.setEnable(true);
        s2sqlParser.setChatModelId(2);
        s2sqlParser.setChatModelConfig(null);
        s2sqlParser.setAppModule(null);

        ChatModelConfig chatModelConfig = new ChatModelConfig("OPEN_AI",
                "https://dashscope.aliyuncs.com/compatible-mode/v1",
                "sk-e395e38b133a44b98142207c96970bb5", "qwen-plus-latest",
                "2024-02-01", 0.0, 60L, null, null, null, 3,
                false, false, false);
        s2sqlParser.setChatModelConfig(chatModelConfig);
        chatAppMap.put("S2SQL_PARSER", s2sqlParser);

        return chatAppMap;
    }

    public void transToSql(String sql) {
//        String sql =
//                "WITH recent_data AS (SELECT 发电量值,数据日期 FROM 关键指标数据集 " +
//                        "WHERE 数据日期 >= '2024-12-01' AND 数据日期 <= '2024-12-15') " +
//                        "SELECT 发电量值 FROM recent_data ORDER BY 数据日期 DESC LIMIT 1000";
        //获取指标维度
//        SemanticTranslateResp explain = semanticLayerService
//                .translate(QueryReqBuilder.buildS2SQLReq(sql, dataSet), User.getDefaultUser());
//        QuerySqlReq querySqlReq = QueryReqBuilder.buildS2SQLReq(sql, dataSet);
//        log.info(querySqlReq.getSql());

        QueryStatement queryStatement = null;
        queryStatement = new QueryStatement();
        queryStatement.setSql(sql);

        //dataSet : 关键指标数据集
        String dsName = "关键指标数据集";
        //metrics : 指标，如：发电量值等,mock these data,use mapper to mock
        List<SchemaElement> metrics = new ArrayList<>();
        metrics = mockMetrics();
        //dimensions : 维度，如：数据日期等,mock these data,use mapper to mock
        List<SchemaElement> dimensions = new ArrayList<>();
        dimensions = mockDimensions();
        //6. tags : 标签，如：发电量值等
        //7. dimensionValues : 维度值，如：2024-12-01等
        //8. terms : 术语，如：发电量值等
        //9. queryConfig : 查询配置，如：where条件，limit等

        //print metrics and dimensions
        log.info("metrics: {}", metrics.stream().map(SchemaElement::getName).collect(Collectors.toList()));
        log.info("dimensions: {}", dimensions.stream().map(SchemaElement::getName).collect(Collectors.toList()));

        //replace the key in sql with bizName
        for (SchemaElement schemaElement : metrics) {
            String name = schemaElement.getName();
            String bizName = schemaElement.getBizName();

            sql = StrUtil.replace(sql, name, bizName);
        }
        //replace the key in sql with bizName,use dimensions
        for (SchemaElement schemaElement : dimensions) {
            String name = schemaElement.getName();
            String bizName = schemaElement.getBizName();

            sql = StrUtil.replace(sql, name, bizName);
        }
        //replace the key in sql with bizName,print the sql
        log.info("1------- sql: {}", sql);
        //先简单处理替换数据表名
        String tableName = "bi_indicator";
        sql = StrUtil.replace(sql, dsName, tableName);

        log.info("2------- sql: {}", sql);


        /**
         * 复杂sql的处理，后续进行研究
         * need to replace the table name in sql with bizName,the bizName is the table name in  dataSet model
         * like: 生产数据集
         */
//        List<String> withNameList = SqlUtil.getWithName(sql);
//        Select selectStatement = SqlUtil.getSelect(sql);
//        List<PlainSelect> plainSelectList = SqlUtil.getWithItem(selectStatement);
//
//        String tableName = "bi_indicator";
//        if (!CollectionUtils.isEmpty(plainSelectList)) {
//            plainSelectList.forEach(
//                    plainSelect -> SqlUtil.processPlainSelect(plainSelect, tableName, withNameList));
//        }


        //增加sql的wehere条件，如果有配置的话
//        if (dataSet.getSqlWhere() != null) {
//            sql = sql + " " + dataSet.getSqlWhere();
//        }


//        assertNotNull(explain);
//        assertNotNull(explain.getQuerySQL());
//        log.info(explain.getQuerySQL());
////        assertTrue(explain.getQuerySQL().toLowerCase().contains("department"));
////        assertTrue(explain.getQuerySQL().toLowerCase().contains("count(1)"));
//        executeSql(explain.getQuerySQL());
    }


    /**
     * mock metrics
     * sjrq:数据日期
     *
     * @return
     */
    private List<SchemaElement> mockDimensions() {
        List<SchemaElement> dimensions = new ArrayList<>();
        List<String> names = ListUtil.of("数据日期");
        List<String> bizNames = ListUtil.of("btime");
        for (int i = 0; i < names.size(); i++) {
            SchemaElement schemaElement = new SchemaElement();
            schemaElement.setName(names.get(i));
            schemaElement.setBizName(bizNames.get(i));
            schemaElement.setId((long) i);
            schemaElement.setType(SchemaElementType.DIMENSION);
            dimensions.add(schemaElement);
        }
        return dimensions;
    }

    /**
     * mock metrics
     * fdlz:发电量值
     * gdmh:供电煤耗
     * sccydl:生产厂用电率
     * zhcydl:综合厂用电率
     * dwfdyh:单位发电油耗
     * fdbsl:发电补水率
     *
     * @return
     */
    private List<SchemaElement> mockMetrics() {
        List<SchemaElement> metrics = new ArrayList<>();
        List<String> names = ListUtil.of("发电量值", "供电煤耗", "生产厂用电率", "综合厂用电率", "单位发电油耗", "发电补水率");
        List<String> bizNames = ListUtil.of("fdlz", "gdmh", "sccydl", "zhcydl", "dwfdyh", "fdbsl");
        for (int i = 0; i < names.size(); i++) {
            SchemaElement schemaElement = new SchemaElement();
            schemaElement.setName(names.get(i));
            schemaElement.setBizName(bizNames.get(i));
            schemaElement.setId((long) i);
            schemaElement.setType(SchemaElementType.METRIC);
            metrics.add(schemaElement);
        }

        return metrics;
    }


}



